# -*- coding: utf-8 -*-
"""
@Time    : 2024/9/3 10:38 
@Author  : ZhangShenao 
@File    : 5.消息的更新与删除.py
@Desc    : 实现消息的更新与删除

LangGraph提供了RemoveMessage类,配合归纳函数add_messages
可以根据消息id,实现消息的更新或删除
其核心思想是归纳函数add_messages底层针对更新的消息类型做了检测
如果检测到是RemoveMessage类型,则不会新增消息,而是将指定id的消息删除
而如果检测到消息不是RemoveMessage类型,而是普通消息,则相同id的消息会进行覆盖,进而实现了更新的操作

删除和更新的操作,本质上都是通过add_messages这个归纳函数来实现的

特别注意:
绝大部分模型期望消息列表存在某些规则
例如:有些模型期望以user消息开头,有些模型期望所有带有工具调用的消息后面都跟着工具消息
删除或修改消息时,需要确保不会违反这些规则
"""
import dotenv
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.messages import RemoveMessage, HumanMessage, AIMessage
from langchain_core.runnables import RunnableConfig
from langgraph.graph import MessagesState, StateGraph

# 加载环境变量
dotenv.load_dotenv()

# 创建LLM
llm = ChatZhipuAI(model="glm-4-plus")


# 定义LLM节点
def llm_node(state: MessagesState, config: RunnableConfig) -> MessagesState:
    """LLM节点"""

    # 从状态中获取消息列表
    msgs = state["messages"]

    # 调用LLM,获取生成内容
    resp = llm.invoke(msgs)

    # 返回更新后的状态
    return {"messages": [resp]}


# 定义创建删除消息节点
def remove_human_msg_node(state: MessagesState, config: RunnableConfig) -> MessagesState:
    """
    删除HumanMessage节点,将消息列表中的第一条HumanMessage删除
    :param state: 消息列表状态
    :param config: 运行时配置
    :return: 更新后的状态
    """

    # 取出状态中的第一条HumanMessage
    human_msg = state["messages"][0]

    # 返回RemoveMessage类型的消息,并且指定要删除的消息id为HumanMessage的id
    remove_msg = RemoveMessage(id=human_msg.id)

    # 返回更新后的状态
    return {"messages": [remove_msg]}


# 定义修改消息的节点
def update_ai_msg_node(state: MessagesState, config: RunnableConfig) -> MessagesState:
    """
    修改AIMessage节点,修改消息列表中的最后一条AIMessage的内容
    :param state: 消息列表状态
    :param config: 运行时配置
    :return: 更新后的状态
    """

    # 从状态数据中,获取最后一条AIMessage
    ai_msg = state["messages"][-1]

    # 创建更新后的AIMessage,消息id与原消息相同
    update_msg = AIMessage(id=ai_msg.id, content="修改后的AI消息：" + ai_msg.content)

    # 返回更新后的状态,LangGraph内部会根据消息ID进行更新
    return {"messages": [update_msg]}


# 创建图构建器,并指定状态为MessagesState类型
graph_builder = StateGraph(MessagesState)

# 添加节点
graph_builder.add_node(node="llm", action=llm_node)
graph_builder.add_node(node="remove_human_msg", action=remove_human_msg_node)
graph_builder.add_node(node="update_ai_msg", action=update_ai_msg_node)

# 添加边
graph_builder.set_entry_point(key="llm")  # 指定开始节点
graph_builder.add_edge(start_key="llm", end_key="remove_human_msg")
graph_builder.add_edge(start_key="remove_human_msg", end_key="update_ai_msg")
graph_builder.set_finish_point(key="update_ai_msg")  # 指定结束节点

# 将图架构编译成Runnable可运行组件
graph = graph_builder.compile()

# 定义初始状态
init_state = {"messages": [HumanMessage(content="你好，我是zsa，我喜欢写代码和打篮球")]}

# 执行图,获取最终状态
final_state = graph.invoke(init_state)

# 最终状态中,消息列表只包含一条修改后的AIMessage
print(final_state)
